package in.ac.iitb.cse.nlp.poswordpredictor;

import in.ac.iitb.cse.nlp.postagger.data.DataMaps;
import in.ac.iitb.cse.nlp.postagger.data.Matrix;
import in.ac.iitb.cse.nlp.postagger.data.Tag;
import in.ac.iitb.cse.nlp.postagger.data.Word;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class InverseEmissionMatrix implements Matrix {

	// Internal map which stores the emission table
	private Map<Tag, HashMap<Word, Integer>> emissionTable = new HashMap<Tag, HashMap<Word, Integer>>();

	public void update(String row, String col) {
		Tag tag = DataMaps.tags.get(row);
		if (tag == null) {
			tag = new Tag(row);
			DataMaps.tags.put(row, tag);
		}
		HashMap<Word, Integer> map = emissionTable.get(tag);
		if (map == null) {
			map = new HashMap<Word, Integer>();
			emissionTable.put(tag, map);
		}
		Word word = DataMaps.words.get(col);
		if (word == null) {
			word = new Word(col);
			DataMaps.words.put(col, word);
		}
		Integer count = map.get(word);
		if (count == null) {
			map.put(word, 1);
		} else {
			map.put(word, count + 1);
		}
	}

	public int getCount(String row, String col) {
		Tag tag = DataMaps.tags.get(row);
		if (tag != null) {
			HashMap<Word, Integer> hashMap = emissionTable.get(tag);
			if (hashMap != null) {
				Word word = DataMaps.words.get(col);
				if (tag != null) {
					Integer integer = hashMap.get(word);
					if (integer != null) {
						return integer;
					}
				}
			}
		}
		return 0;
	}

	/**
	 * Usage: getProbability(tag, word) = Prob(word/tag)
	 */
	public double getProbability(String row, String col) {
		int countOfColGivenRow = getCount(row, col);
		if (countOfColGivenRow != 0) {
			Tag tag = DataMaps.tags.get(row);
			return ((double) countOfColGivenRow / (double) tag.getCount());
		}
		return 0;
	}

	List<Prediction> getSuitablePrediction(Prediction prediction, double factor) {
		Tag tag = DataMaps.tags.get(prediction.word);
		List<Prediction> predictions = new ArrayList<Prediction>();
		if (tag != null) {
			HashMap<Word, Integer> hashMap = emissionTable.get(tag);
			List<Prediction> temp = new ArrayList<Prediction>();
			if (hashMap != null) {
				for (Word key : hashMap.keySet()) {
					temp.add(new Prediction(key.getString(), getProbability(
							tag.getString(), key.getString())));
				}
			}
			Collections.sort(temp);
			for (int i = 0; i < temp.size()
					&& i < Constants.MAX_NUM_OF_PREDICTONS * factor; i++) {
				// Mulitply each prediction by its weight
				predictions.add(new Prediction(temp.get(i).word,
						temp.get(i).probability * prediction.probability));
			}
		}
		System.out.println(factor+","+prediction.word + ":" + predictions);
		return predictions;
	}

	@Override
	public String toString() {
		StringBuffer buffer = new StringBuffer();
		for (Tag tag : emissionTable.keySet()) {
			for (Word word : emissionTable.get(tag).keySet()) {
				buffer.append("(" + word + "," + tag + ")="
						+ emissionTable.get(tag).get(word) + ":"
						+ getProbability(tag.toString(), word.toString()));
			}
			buffer.append("\n");
		}
		return buffer.toString();
	}

}
